package sshex

import (
	"bufio"
	"context"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"net"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"gitee.com/shanyou/sshplus"
	"gitee.com/shanyou/sshplus/rlog"
	"gitee.com/shanyou/sshplus/utils"
	"github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"golang.org/x/crypto/ssh/knownhosts"
	"golang.org/x/term"
)

const (
	defaultTermName string = "xterm"
	maxPromptTries  int    = 3
)

type ClientAuth struct {
	User            string
	RemoteAddr      string
	NewAgent        bool
	IgnoreHostKey   bool
	Ciphers         []string
	KnownHostsFiles []string
	MaxPromptTries  int
	Logger          logrus.FieldLogger
	LocalAgent      agent.Agent
	certChecker     *ssh.CertChecker
	ctx             context.Context
	cancel          context.CancelFunc
}

func NewClientAuth(ctx context.Context) *ClientAuth {
	c, cancel := context.WithCancel(ctx)
	localAgent := agent.NewKeyring()
	l := rlog.Logger.WithField("com", "ssh-client")
	auth := &ClientAuth{
		ctx:        c,
		cancel:     cancel,
		LocalAgent: localAgent,
		Logger:     l,
	}

	return auth
}

func (auth *ClientAuth) AddKeyWithBytes(keyBytes, certBytes []byte) error {
	block, _ := pem.Decode(keyBytes)
	key, err := x509.ParseECPrivateKey(block.Bytes)
	if err != nil {
		return fmt.Errorf("add key to agent error %v", err)
	}
	pubKey, _, _, _, err := ssh.ParseAuthorizedKey(certBytes)
	if err != nil {
		return fmt.Errorf("add key to agent error %v", err)
	}
	cert, ok := pubKey.(*ssh.Certificate)
	if !ok {
		return fmt.Errorf("cert file format error")
	}
	// add sshkey
	err = auth.LocalAgent.Add(agent.AddedKey{
		PrivateKey:  key,
		Certificate: cert,
	})

	if err != nil {
		return fmt.Errorf("add key to agent error %v", err)
	}

	return nil
}

func (auth *ClientAuth) Close() {
	auth.cancel()
}

func (auth *ClientAuth) BuildConfig() *ssh.ClientConfig {
	var sshAgent agent.Agent
	agentSocket := os.Getenv("SSH_AUTH_SOCK")
	if len(agentSocket) > 0 && !auth.NewAgent {
		uconn, err := net.Dial("unix", agentSocket)
		if err != nil {
			auth.Logger.Warningf("dail ssh agent sock error %v", err)
		} else {
			sshAgent = agent.NewClient(uconn)
			go func() {
				<-auth.ctx.Done()
				uconn.Close()
			}()
		}
	} else {
		sshAgent = auth.LocalAgent
	}
	config := &ssh.ClientConfig{
		User: auth.User,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeysCallback(sshAgent.Signers),
			ssh.RetryableAuthMethod(ssh.KeyboardInteractive(auth.keyboardInteractiveChallenge), maxPromptTries),
			ssh.RetryableAuthMethod(ssh.PasswordCallback(auth.PasswordHandler), maxPromptTries),
		},
		ClientVersion: "SSH-2.0-" + sshplus.ProtocolVersion,
	}

	config.HostKeyCallback = auth.HostKeyCallback
	return config
}

func (auth *ClientAuth) SetHostCAPubWithBytes(certBytes []byte) error {
	pubKey, _, _, _, err := ssh.ParseAuthorizedKey(certBytes)
	if err != nil {
		return err
	}

	var certChecker ssh.CertChecker
	certChecker.IsHostAuthority = func(auth ssh.PublicKey, address string) bool {
		return utils.KeysEqual(pubKey, auth)
	}
	auth.certChecker = &certChecker
	return nil
}

func (auth *ClientAuth) HostKeyCallback(hostname string, remote net.Addr, key ssh.PublicKey) error {
	if auth.IgnoreHostKey {
		return nil
	}

	var hname string
	var raddr net.Addr
	if remote.Network() == "unix" {
		hname = "localhost:22"
		raddr, _ = net.ResolveTCPAddr("tcp", "localhost:22")
	} else {
		hname = hostname
		raddr = remote
	}
	if auth.certChecker != nil {
		return auth.certChecker.CheckHostKey(hname, raddr, key)
	}

	if len(auth.KnownHostsFiles) == 0 {
		return fmt.Errorf("there is no knownhosts file")
	}

	hostKeyCallback, err := knownhosts.New(auth.KnownHostsFiles...)
	if err != nil {
		return fmt.Errorf("failed to load knownhosts files: %s", err)
	}

	err = hostKeyCallback(hname, raddr, key)

	if err == nil {
		return nil
	}

	_, isCert := key.(*ssh.Certificate)
	keyErr, ok := err.(*knownhosts.KeyError)
	if isCert {
		if answer, err := askAddingUnknownHostKey(hname, raddr, key); err != nil || !answer {
			msg := "host key verification failed"
			if err != nil {
				msg += ": " + err.Error()
			}
			return fmt.Errorf(msg)
		}

		addHostKey(hname, remote, key, auth.KnownHostsFiles[0])
		return nil
	}
	if !ok || len(keyErr.Want) > 0 {
		return err
	}

	if answer, err := askAddingUnknownHostKey(hname, raddr, key); err != nil || !answer {
		msg := "host key verification failed"
		if err != nil {
			msg += ": " + err.Error()
		}
		return fmt.Errorf(msg)
	}

	addHostKey(hname, raddr, key, auth.KnownHostsFiles[0])
	return nil
}

func (auth *ClientAuth) PasswordHandler() (secret string, err error) {
	ttyin, ttyout, err := openTTY()
	if err != nil {
		return "", fmt.Errorf("failed to open tty: %s", err)
	}
	defer closeTTY(ttyin, ttyout)
	//root@192.168.1.1's password:
	prompt := fmt.Sprintf("%s@%s's password: ", auth.User, auth.RemoteAddr)
	return readPassword(ttyin, ttyout, prompt)
}

func (auth *ClientAuth) keyboardInteractiveChallenge(user, instruction string, questions []string, echos []bool) (answers []string, err error) {
	ttyin, ttyout, err := openTTY()
	if err != nil {
		return answers, fmt.Errorf("failed to open tty: %s", err)
	}
	defer closeTTY(ttyin, ttyout)
	answers = make([]string, len(questions))
	var strs []string
	if len(questions) > 0 {
		if user != "" {
			strs = append(strs, user)
		}
		if instruction != "" {
			strs = append(strs, instruction)
		}
		if len(strs) > 0 {
			fmt.Fprintln(ttyout, strings.Join(strs, " "))
		} else {
			fmt.Fprintf(ttyout, "Keyboard interactive challenge for %s@%s\n", auth.User, auth.RemoteAddr)
		}
	}
	for i, q := range questions {
		res, err := readPassword(ttyin, ttyout, q)
		if err != nil {
			return answers, err
		}
		answers[i] = res
	}
	return answers, err
}

func readPassword(ttyin, ttyout *os.File, prompt string) (password string, err error) {
	state, err := term.GetState(int(ttyin.Fd()))
	if err != nil {
		return "", fmt.Errorf("failed to get terminal state: %s", err)
	}

	stopC := make(chan struct{})
	defer func() {
		close(stopC)
	}()

	go func() {
		sigC := make(chan os.Signal, 1)
		signal.Notify(sigC, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
		select {
		case <-sigC:
			term.Restore(int(ttyin.Fd()), state)
			os.Exit(1)
		case <-stopC:
		}
	}()

	if prompt == "" {
		fmt.Fprint(ttyout, "Password: ")
	} else {
		fmt.Fprint(ttyout, prompt)
	}

	b, err := term.ReadPassword(int(ttyin.Fd()))
	if err != nil {
		return "", fmt.Errorf("failed to read password: %s", err)
	}

	fmt.Fprint(ttyout, "\n")

	return string(b), nil
}

func askAddingUnknownHostKey(address string, remote net.Addr, key ssh.PublicKey) (bool, error) {
	stopC := make(chan struct{})
	defer func() {
		close(stopC)
	}()

	go func() {
		sigC := make(chan os.Signal, 1)
		signal.Notify(sigC, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
		select {
		case <-sigC:
			os.Exit(1)
		case <-stopC:
		}
	}()

	fmt.Printf("The authenticity of host '%s (%s)' can't be established.\n", address, remote.String())
	fmt.Printf("key fingerprint is %s\n", ssh.FingerprintSHA256(key))
	fmt.Printf("Are you sure you want to continue connecting (yes/no)? ")

	b := bufio.NewReader(os.Stdin)
	for {
		answer, err := b.ReadString('\n')
		if err != nil {
			return false, fmt.Errorf("failed to read answer: %s", err)
		}
		answer = string(strings.ToLower(strings.TrimSpace(answer)))
		if answer == "yes" {
			return true, nil
		} else if answer == "no" {
			return false, nil
		}
		fmt.Print("Please type 'yes' or 'no': ")
	}
}

func addHostKey(hostname string, remote net.Addr, key ssh.PublicKey, knownhostsPath string) {
	normAddr := knownhosts.Normalize(hostname)
	certKey, isCert := key.(*ssh.Certificate)
	go func() {
		// add host to known_hosts
		encoded := knownhosts.HashHostname(normAddr)
		var line string
		if isCert {
			line = "@cert-authority" + "\t" + encoded + "\t" + key.Type() + "\t" + base64.StdEncoding.EncodeToString(certKey.SignatureKey.Marshal())
		} else {
			line = encoded + "\t" + key.Type() + "\t" + base64.StdEncoding.EncodeToString(key.Marshal())
		}

		fmt.Printf("add \"%s\" to known hosts\n", line)

		// append file
		f, ferr := os.OpenFile(knownhostsPath, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
		if ferr != nil {
			rlog.Warningf("append known hosts error %v\n", ferr)
			return
		}
		defer f.Close()
		if _, ferr = f.WriteString(line + "\n"); ferr != nil {
			rlog.Warningf("append known hosts error %v\n", ferr)
		}
	}()
}
